import os
import argparse

def print_directory_tree(root_dir, max_depth=None, prefix=''):
    """递归打印目录树结构"""
    if max_depth is not None and max_depth < 0:
        return
    
    # 获取所有子目录并排序
    entries = []
    try:
        entries = sorted(os.listdir(root_dir))
    except PermissionError:
        return
    
    dirs = [entry for entry in entries if os.path.isdir(os.path.join(root_dir, entry))]
    
    for i, dir_name in enumerate(dirs):
        path = os.path.join(root_dir, dir_name)
        is_last = i == len(dirs) - 1
        
        print(prefix + ('└─' if is_last else '├─') + dir_name)
        
        new_prefix = prefix + ('  ' if is_last else '│ ')
        next_depth = max_depth - 1 if max_depth is not None else None
        print_directory_tree(path, next_depth, new_prefix)

def main():
    parser = argparse.ArgumentParser(description='打印目录树结构')
    parser.add_argument('directory', nargs='?', default='.', help='目标目录（默认为当前目录）')
    parser.add_argument('-d', '--depth', type=int, help='最大递归深度')
    
    args = parser.parse_args()
    
    root_dir = os.path.abspath(args.directory)
    if not os.path.isdir(root_dir):
        print(f"错误: '{root_dir}' 不是有效的目录")
        return
    
    # 显示当前目录名
    print(os.path.basename(root_dir))
    print_directory_tree(root_dir, args.depth)

if __name__ == '__main__':
    main()
